package com.light.dp.flink.sink;

import com.alibaba.fastjson.JSONObject;
import com.clickhouse.jdbc.ClickHouseDataSource;
import com.light.dp.flink.config.dto.ColumnMapping;
import com.light.dp.flink.config.dto.JDBCDatabaseConfig;
import com.light.dp.flink.config.dto.TableMapping;
import com.light.dp.flink.dto.DataChangeInfo;
import org.apache.commons.collections.CollectionUtils;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.functions.sink.RichSinkFunction;
import org.apache.flink.streaming.api.functions.sink.SinkFunction;
import org.apache.flink.streaming.api.functions.timestamps.BoundedOutOfOrdernessTimestampExtractor;
import org.apache.flink.streaming.api.windowing.time.Time;

import java.sql.*;
import java.util.*;

public class ClickHouseIncrementalSink extends RichSinkFunction<DataChangeInfo> {
    private static final int MAX_RETRIES = 3;
    private static final int BATCH_SIZE = 1000;
    private static final long BATCH_INTERVAL = 5000; // 1 秒

    private Connection connection;
    private List<PreparedStatement> insertStatements;
    private List<PreparedStatement> updateStatements;
    private List<PreparedStatement> deleteStatements;
    private JDBCDatabaseConfig targetConfig;
    private List<TableMapping> tableMappings;

    private List<DataChangeInfo> buffer;
    private long lastFlushTime;

    public ClickHouseIncrementalSink(JDBCDatabaseConfig targetConfig, List<TableMapping> tableMappings) {
        this.targetConfig = targetConfig;
        this.tableMappings = tableMappings;
        this.buffer = new ArrayList<>();
        this.lastFlushTime = System.currentTimeMillis();
    }

    @Override
    public void open(Configuration parameters) throws Exception {
        super.open(parameters);
        String url = String.format("jdbc:%s://%s:%d/%s", targetConfig.getJdbcDriver().getJdbcDrive(), targetConfig.getHostname(), targetConfig.getPort(), targetConfig.getDatabase());
        ClickHouseDataSource dataSource = new ClickHouseDataSource(url);
        connection = dataSource.getConnection(targetConfig.getUsername(), targetConfig.getPassword());
        insertStatements = new ArrayList<>();
        updateStatements = new ArrayList<>();
        deleteStatements = new ArrayList<>();

        for (TableMapping tableMapping : tableMappings) {
            List<ColumnMapping> columnMappings = tableMapping.getColumnMappings();
            if (CollectionUtils.isEmpty(columnMappings)) {
                String targetTable = tableMapping.getTargetTable();
                List<String> columns = getTableColumns(targetTable);
                List<String> primaryKeys = getTablePrimaryKeys(targetTable);

                columnMappings = new ArrayList<>();
                for (String column : columns) {
                    boolean isPrimaryKey = primaryKeys.contains(column);
                    columnMappings.add(new ColumnMapping(column, column, isPrimaryKey));
                }
                tableMapping.setColumnMappings(columnMappings);
            }

            StringBuilder insertColumns = new StringBuilder();
            StringBuilder insertValues = new StringBuilder();
            StringBuilder updateClause = new StringBuilder();
            StringBuilder deleteCondition = new StringBuilder();
            List<String> primaryKeys = new ArrayList<>();

            for (int i = 0; i < columnMappings.size(); i++) {
                ColumnMapping mapping = columnMappings.get(i);
                String targetColumn = mapping.getTargetColumn();
                if (i > 0) {
                    insertColumns.append(", ");
                    insertValues.append(", ");
                }
                insertColumns.append(targetColumn);
                insertValues.append("?");

                if (mapping.getIsTargetPrimaryKey()) {
                    primaryKeys.add(targetColumn);
                    if (deleteCondition.length() > 0) {
                        deleteCondition.append(" AND ");
                    }
                    deleteCondition.append(targetColumn).append(" = ?");
                } else {
                    if (updateClause.length() > 0) {
                        updateClause.append(", ");
                    }
                    updateClause.append(targetColumn).append(" = ?");
                }
            }

            // 构建插入语句
            String insertQuery = "INSERT INTO " + tableMapping.getTargetTable() + " (" + insertColumns + ") VALUES (" + insertValues + ")";
            insertStatements.add(connection.prepareStatement(insertQuery));

            // 构建更新语句
            StringBuilder updateQuery = new StringBuilder("ALTER TABLE " + tableMapping.getTargetTable() + " UPDATE " + updateClause);
            updateQuery.append(" WHERE ");
            for (int i = 0; i < primaryKeys.size(); i++) {
                if (i > 0) {
                    updateQuery.append(" AND ");
                }
                updateQuery.append(primaryKeys.get(i)).append(" = ?");
            }
            updateStatements.add(connection.prepareStatement(updateQuery.toString()));

            // 构建删除语句
            String deleteQuery = "ALTER TABLE " + tableMapping.getTargetTable() + " DELETE WHERE " + deleteCondition;
            deleteStatements.add(connection.prepareStatement(deleteQuery));
        }
    }

    @Override
    public void invoke(DataChangeInfo value, SinkFunction.Context context) throws Exception {
        if (value.getEventType() != "CREATE" && value.getEventType() != "UPDATE" && value.getEventType() != "DELETE"){
            return;
        }
        buffer.add(value);
        long currentTime = System.currentTimeMillis();
        if (buffer.size() >= BATCH_SIZE || currentTime - lastFlushTime >= BATCH_INTERVAL) {
            flushBuffer();
            lastFlushTime = currentTime;
        }
    }

    private void flushBuffer() throws Exception {
        Map<String, List<DataChangeInfo>> operationMap = new HashMap<>();
        operationMap.put("CREATE", new ArrayList<>());
        operationMap.put("UPDATE", new ArrayList<>());
        operationMap.put("DELETE", new ArrayList<>());

        for (DataChangeInfo info : buffer) {
            if (operationMap.get(info.getEventType()) == null){
                continue;
            }
            operationMap.get(info.getEventType()).add(info);
        }

        for (Map.Entry<String, List<DataChangeInfo>> entry : operationMap.entrySet()) {
            String operation = entry.getKey();
            List<DataChangeInfo> infos = entry.getValue();
            switch (operation) {
                case "CREATE":
                    for (int i = 0; i < tableMappings.size(); i++) {
                        TableMapping tableMapping = tableMappings.get(i);
                        List<ColumnMapping> columnMappings = tableMapping.getColumnMappings();
                        PreparedStatement insertStatement = insertStatements.get(i);
                        for (DataChangeInfo info : infos) {
                            handleInsert(JSONObject.parseObject(info.getAfterData()), columnMappings, insertStatement);
                        }
                        insertStatement.executeBatch();
                        insertStatement.clearBatch();
                    }
                    break;
                case "UPDATE":
                    for (int i = 0; i < tableMappings.size(); i++) {
                        TableMapping tableMapping = tableMappings.get(i);
                        List<ColumnMapping> columnMappings = tableMapping.getColumnMappings();
                        List<String> primaryKeys = new ArrayList<>();
                        for (ColumnMapping mapping : columnMappings) {
                            if (mapping.getIsTargetPrimaryKey()) {
                                primaryKeys.add(mapping.getTargetColumn());
                            }
                        }
                        PreparedStatement updateStatement = updateStatements.get(i);
                        for (DataChangeInfo info : infos) {
                            handleUpdate(JSONObject.parseObject(info.getAfterData()), columnMappings, primaryKeys, updateStatement);
                        }
                        updateStatement.executeBatch();
                        updateStatement.clearBatch();
                    }
                    break;
                case "DELETE":
                    for (int i = 0; i < tableMappings.size(); i++) {
                        TableMapping tableMapping = tableMappings.get(i);
                        List<String> primaryKeys = new ArrayList<>();
                        List<ColumnMapping> columnMappings = tableMapping.getColumnMappings();
                        for (ColumnMapping mapping : columnMappings) {
                            if (mapping.getIsTargetPrimaryKey()) {
                                primaryKeys.add(mapping.getTargetColumn());
                            }
                        }
                        PreparedStatement deleteStatement = deleteStatements.get(i);
                        for (DataChangeInfo info : infos) {
                            handleDelete(JSONObject.parseObject(info.getBeforeData()), primaryKeys, deleteStatement);
                        }
                        deleteStatement.executeBatch();
                        deleteStatement.clearBatch();
                    }
                    break;
                default:
                    System.out.println("Unknown operation type: " + operation);
            }
        }
        buffer.clear();
    }

    private void handleInsert(JSONObject json, List<ColumnMapping> columnMappings, PreparedStatement statement) throws SQLException {
        int retries = 0;
        while (retries < MAX_RETRIES) {
            try {
                for (int j = 0; j < columnMappings.size(); j++) {
                    ColumnMapping mapping = columnMappings.get(j);
                    String sourceColumn = mapping.getSourceColumn();
                    Object val = json.get(sourceColumn);
                    statement.setObject(j + 1, val);
                }
                statement.addBatch();
                break;
            } catch (SQLException e) {
                retries++;
                if (retries == MAX_RETRIES) {
                    throw e;
                }
                try {
                    Thread.sleep(1000);
                } catch (InterruptedException ie) {
                    Thread.currentThread().interrupt();
                }
            }
        }
    }

    private void handleUpdate(JSONObject json, List<ColumnMapping> columnMappings, List<String> primaryKeys, PreparedStatement statement) throws SQLException {
        int retries = 0;
        while (retries < MAX_RETRIES) {
            try {
                int paramIndex = 1;
                // 设置更新字段的值
                for (ColumnMapping mapping : columnMappings) {
                    if (!mapping.getIsTargetPrimaryKey()) {
                        String sourceColumn = mapping.getSourceColumn();
                        Object val = json.get(sourceColumn);
                        statement.setObject(paramIndex++, val);
                    }
                }
                // 设置主键条件的值
                for (String primaryKey : primaryKeys) {
                    for (ColumnMapping mapping : columnMappings) {
                        if (mapping.getTargetColumn().equals(primaryKey)) {
                            String sourceColumn = mapping.getSourceColumn();
                            Object val = json.get(sourceColumn);
                            statement.setObject(paramIndex++, val);
                            break;
                        }
                    }
                }
                statement.addBatch();
                break;
            } catch (SQLException e) {
                retries++;
                if (retries == MAX_RETRIES) {
                    throw e;
                }
                try {
                    Thread.sleep(1000);
                } catch (InterruptedException ie) {
                    Thread.currentThread().interrupt();
                }
            }
        }
    }

    private void handleDelete(JSONObject json, List<String> primaryKeys, PreparedStatement statement) throws SQLException {
        int retries = 0;
        while (retries < MAX_RETRIES) {
            try {
                int paramIndex = 1;
                for (String primaryKey : primaryKeys) {
                    for (ColumnMapping mapping : tableMappings.get(0).getColumnMappings()) {
                        if (mapping.getTargetColumn().equals(primaryKey)) {
                            String sourceColumn = mapping.getSourceColumn();
                            Object val = json.get(sourceColumn);
                            statement.setObject(paramIndex++, val);
                            break;
                        }
                    }
                }
                statement.addBatch();
                break;
            } catch (SQLException e) {
                retries++;
                if (retries == MAX_RETRIES) {
                    throw e;
                }
                try {
                    Thread.sleep(1000);
                } catch (InterruptedException ie) {
                    Thread.currentThread().interrupt();
                }
            }
        }
    }

    @Override
    public void close() throws Exception {
        if (!buffer.isEmpty()) {
            flushBuffer();
        }
        closeStatements(insertStatements);
        closeStatements(updateStatements);
        closeStatements(deleteStatements);
        if (connection != null) {
            connection.close();
        }
    }

    private void closeStatements(List<PreparedStatement> statements) throws SQLException {
        for (PreparedStatement statement : statements) {
            if (statement != null) {
                statement.close();
            }
        }
    }

    private List<String> getTableColumns(String tableName) throws SQLException {
        List<String> columns = new ArrayList<>();
        Statement statement = connection.createStatement();
        ResultSet resultSet = statement.executeQuery("SELECT * FROM " + tableName + " LIMIT 0");
        int columnCount = resultSet.getMetaData().getColumnCount();
        for (int i = 1; i <= columnCount; i++) {
            columns.add(resultSet.getMetaData().getColumnName(i));
        }
        resultSet.close();
        statement.close();
        return columns;
    }

    private List<String> getTablePrimaryKeys(String tableName) throws SQLException {
        List<String> primaryKeys = new ArrayList<>();
        String query = "SELECT name " +
                "FROM system.columns " +
                "WHERE database = ? " +
                "AND table = ? " +
                "AND is_in_primary_key = 1";
        PreparedStatement preparedStatement = connection.prepareStatement(query);
        preparedStatement.setString(1, targetConfig.getDatabase());
        preparedStatement.setString(2, tableName);
        ResultSet resultSet = preparedStatement.executeQuery();
        while (resultSet.next()) {
            primaryKeys.add(resultSet.getString("name"));
        }
        resultSet.close();
        preparedStatement.close();
        return primaryKeys;
    }
}